package tsj.ai.alibaba.data.analysis.prompt;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.BooleanUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.prompt.PromptTemplate;
import tsj.ai.alibaba.data.analysis.common.connector.config.DbConfig;
import tsj.ai.alibaba.data.analysis.dto.BusinessKnowledgeDTO;
import tsj.ai.alibaba.data.analysis.dto.schema.ColumnDTO;
import tsj.ai.alibaba.data.analysis.dto.schema.SchemaDTO;
import tsj.ai.alibaba.data.analysis.dto.schema.TableDTO;
import tsj.ai.alibaba.data.analysis.entity.SemanticModel;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.stream.Collectors;
import tsj.ai.alibaba.data.analysis.common.enums.BizDataSourceTypeEnum;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-13 14:17:35
 */
@Slf4j
public class PromptHelper {

    private static final List<String> DATE_TIME_TYPES = Arrays.asList("DATE", "TIME", "DATETIME", "TIMESTAMP");

    /**
     * 构建时间转换提示词
     * @param query 用户查询
     * @return 时间转换提示词
     */
    public static String buildTimeConversionPrompt(String query) {
        Map<String, Object> promptMap = new HashMap<>();
        promptMap.put("current_time_info",
                LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        promptMap.put("query", query);

        PromptTemplate promptTemplate = PromptConstant.getTimeConversionPromptTemplate();
        return promptTemplate.render(promptMap);
    }

    public static String buildQueryToKeywordsPrompt(String question) {
        Map<String, Object> params = new HashMap<>();
        params.put("question", question);
        return PromptConstant.getQuestionToKeywordsPromptTemplate().render(params);
    }

    public static String buildBusinessKnowledgePrompt(List<BusinessKnowledgeDTO> businessKnowledgeDTOS) {
        Map<String, Object> params = new HashMap<>();
        String businessKnowledge = CollectionUtils.isEmpty(businessKnowledgeDTOS) ? ""
                : StringUtils.join(businessKnowledgeDTOS, ";\n");
        params.put("businessKnowledge", businessKnowledge);
        return PromptConstant.getBusinessKnowledgePromptTemplate().render(params);
    }

    public static String buildSemanticModelPrompt(List<SemanticModel> semanticModels) {
        Map<String, Object> params = new HashMap<>();
        String semanticModel = CollectionUtils.isEmpty(semanticModels) ? ""
                : semanticModels.stream().map(SemanticModel::getPromptInfo).collect(Collectors.joining(";\n"));
        params.put("semanticModel", semanticModel);
        return PromptConstant.getSemanticModelPromptTemplate().render(params);
    }

    public static String buildSemanticConsistenPrompt(String nlReq, String sql) {
        Map<String, Object> params = new HashMap<>();
        params.put("nl_req", nlReq);
        params.put("sql", sql);
        return PromptConstant.getSemanticConsistencyPromptTemplate().render(params);
    }

    public static String buildSqlErrorFixerPrompt(String question, DbConfig dbConfig, SchemaDTO schemaDTO,
                                                  List<String> evidenceList, String errorSql, String errorMessage) {
        String evidence = StringUtils.join(evidenceList, ";\n");
        String schemaInfo = buildMixMacSqlDbPrompt(schemaDTO, true);
        String dialect = BizDataSourceTypeEnum.fromTypeName(dbConfig.getDialectType()).getDialect();

        Map<String, Object> params = new HashMap<>();
        params.put("dialect", dialect);
        params.put("question", question);
        params.put("schema_info", schemaInfo);
        params.put("evidence", evidence);
        params.put("error_sql", errorSql);
        params.put("error_message", errorMessage);

        return PromptConstant.getSqlErrorFixerPromptTemplate().render(params);
    }

    public static String buildMixMacSqlDbPrompt(SchemaDTO schemaDTO, Boolean withColumnType) {
        StringBuilder sb = new StringBuilder();
        sb.append("【DB_ID】 ").append(schemaDTO.getName() == null ? "" : schemaDTO.getName()).append("\n");
        for (TableDTO tableDTO : schemaDTO.getTable()) {
            sb.append(buildMixMacSqlTablePrompt(tableDTO, withColumnType)).append("\n");
        }
        if (CollectionUtils.isNotEmpty(schemaDTO.getForeignKeys())
                && CollectionUtils.isNotEmpty(schemaDTO.getForeignKeys().get(0))) {
            sb.append("【Foreign keys】\n").append(StringUtils.join(schemaDTO.getForeignKeys().get(0), "\n"));
        }
        return sb.toString();
    }

    public static String buildMixMacSqlTablePrompt(TableDTO tableDTO, Boolean withColumnType) {
        StringBuilder sb = new StringBuilder();
        // sb.append("# Table:
        // ").append(tableDTO.getName()).append(StringUtils.isBlank(tableDTO.getDescription())
        // ? "" : ", " + tableDTO.getDescription()).append("\n");
        sb.append("# Table: ").append(tableDTO.getName());
        if (!StringUtils.equals(tableDTO.getName(), tableDTO.getDescription())) {
            sb.append(StringUtils.isBlank(tableDTO.getDescription()) ? "" : ", " + tableDTO.getDescription())
                    .append("\n");
        }
        else {
            sb.append("\n");
        }
        sb.append("[\n");
        List<String> columnLines = new ArrayList<>();
        for (ColumnDTO columnDTO : tableDTO.getColumn()) {
            StringBuilder line = new StringBuilder();
            line.append("(")
                    .append(columnDTO.getName())
                    .append(BooleanUtils.isTrue(withColumnType)
                            ? ":" + StringUtils.defaultString(columnDTO.getType(), "").toUpperCase(Locale.ROOT) : "");
            if (!StringUtils.equals(columnDTO.getDescription(), columnDTO.getName())) {
                line.append(", ").append(StringUtils.defaultString(columnDTO.getDescription(), ""));
            }
            if (CollectionUtils.isNotEmpty(tableDTO.getPrimaryKeys())
                    && tableDTO.getPrimaryKeys().contains(columnDTO.getName())) {
                line.append(", Primary Key");
            }
            List<String> enumData = Optional.ofNullable(columnDTO.getData())
                    .orElse(new ArrayList<>())
                    .stream()
                    .filter(d -> !StringUtils.isEmpty(d))
                    .collect(Collectors.toList());
            if (CollectionUtils.isNotEmpty(enumData) && !"id".equals(columnDTO.getName())) {
                line.append(", Examples: [");
                List<String> data = new ArrayList<>(enumData.subList(0, Math.min(3, enumData.size())));
                line.append(StringUtils.join(data, ",")).append("]");
            }
            else if (CollectionUtils.isNotEmpty(columnDTO.getSamples())) {
                List<String> data = columnDTO.getSamples().subList(0, Math.min(3, columnDTO.getSamples().size()));
                data = data.stream().filter(item -> StringUtils.isNotBlank(item)).collect(Collectors.toList());
                if (CollectionUtils.isNotEmpty(data)) {
                    line.append(", Examples: [");
                    data = processSamples(data, columnDTO);
                    line.append(StringUtils.join(data, ",")).append("]");
                }
            }
            line.append(")");
            columnLines.add(line.toString());
        }
        sb.append(StringUtils.join(columnLines, ",\n"));
        sb.append("\n]");
        return sb.toString();
    }

    private static List<String> processSamples(List<String> samples, ColumnDTO columnDTO) {
        final List<String> data = new ArrayList<>(samples);
        if (data.stream().anyMatch(item -> item.length() > 50)) {
            return new ArrayList<>();
        }
        String type = columnDTO.getType();
        if (type != null && DATE_TIME_TYPES.contains(type.toUpperCase(Locale.ROOT))) {
            return data.isEmpty() ? Collections.emptyList() : Collections.singletonList(data.get(0));
        }
        if (type != null && type.equalsIgnoreCase("NUMBER")) {
            return data.isEmpty() ? Collections.emptyList() : Collections.singletonList(data.get(0));
        }
        String columnName = columnDTO.getName();
        if (columnName != null && columnName.trim().toLowerCase(Locale.ROOT).endsWith("id")) {
            return data.isEmpty() ? Collections.emptyList() : Collections.singletonList(data.get(0));
        }
        List<String> longSamples = data.stream().filter(item -> item.length() > 20).collect(Collectors.toList());
        if (CollectionUtils.isNotEmpty(longSamples)) {
            return Collections.singletonList(longSamples.get(0));
        }

        return data;
    }

    public static List<String> buildMixSqlGeneratorPrompt(String question, DbConfig dbConfig, SchemaDTO schemaDTO,
                                                          List<String> evidenceList) {
        String evidence = StringUtils.join(evidenceList, ";\n");
        String schemaInfo = buildMixMacSqlDbPrompt(schemaDTO, true);
        String dialect = BizDataSourceTypeEnum.fromTypeName(dbConfig.getDialectType()).getDialect();
        Map<String, Object> params = new HashMap<>();
        params.put("dialect", dialect);
        params.put("question", question);
        params.put("schema_info", schemaInfo);
        params.put("evidence", evidence);
        List<String> prompts = new ArrayList<>();
        prompts.add(PromptConstant.getMixSqlGeneratorSystemPromptTemplate().render(params));
        prompts.add(PromptConstant.getMixSqlGeneratorPromptTemplate().render(params));
        return prompts;
    }
    public static String buildMixSelectorPrompt(List<String> evidences, String question, SchemaDTO schemaDTO) {
        String schemaInfo = buildMixMacSqlDbPrompt(schemaDTO, true);
        Map<String, Object> params = new HashMap<>();
        params.put("schema_info", schemaInfo);
        params.put("question", question);
        String evidence = CollectionUtils.isEmpty(evidences) ? "" : StringUtils.join(evidences, ";\n");
        params.put("evidence", evidence);
        return PromptConstant.getMixSelectorPromptTemplate().render(params);
    }
}
